package org.ktx.osr.audio

import android.content.Context
import android.media.MediaCodec
import android.media.MediaExtractor
import android.media.MediaFormat
import android.media.MediaMuxer
import android.util.Log
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.withContext
import java.io.File
import java.io.IOException
import java.nio.ByteBuffer

/**
 * 媒体合并器
 * 用于将音频和视频文件合并为最终的录制文件
 */
class MediaMerger(private val context: Context) {
    
    companion object {
        private const val TAG = "MediaMerger"
    }
    
    /**
     * 合并音频和视频文件
     * @param videoFile 视频文件
     * @param audioFile 音频文件
     * @param outputFile 输出文件
     * @return 是否合并成功
     */
    suspend fun mergeAudioVideo(
        videoFile: File,
        audioFile: File,
        outputFile: File
    ): Boolean = withContext(Dispatchers.IO) {
        try {
            Log.d(TAG, "Starting audio-video merge")
            Log.d(TAG, "Video file: ${videoFile.absolutePath}")
            Log.d(TAG, "Audio file: ${audioFile.absolutePath}")
            Log.d(TAG, "Output file: ${outputFile.absolutePath}")
            
            // 检查音频文件格式
            val audioExtractor = MediaExtractor()
            audioExtractor.setDataSource(audioFile.absolutePath)
            val audioTrackIndex = getAudioTrackIndex(audioExtractor)
            val audioFormat = audioExtractor.getTrackFormat(audioTrackIndex)
            val audioMime = audioFormat.getString(MediaFormat.KEY_MIME)
            
            Log.d(TAG, "Audio format: $audioMime")
            Log.d(TAG, "Audio sample rate: ${audioFormat.getInteger(MediaFormat.KEY_SAMPLE_RATE)}")
            Log.d(TAG, "Audio channel count: ${audioFormat.getInteger(MediaFormat.KEY_CHANNEL_COUNT)}")
            
            // 检查音频格式是否有效
            val sampleRate = audioFormat.getInteger(MediaFormat.KEY_SAMPLE_RATE)
            val channelCount = audioFormat.getInteger(MediaFormat.KEY_CHANNEL_COUNT)
            
            val processedAudioFile = when {
                // WAV格式需要转换
                audioMime == "audio/wav" || audioFile.extension.lowercase() == "wav" -> {
                    Log.d(TAG, "Converting WAV to AAC format")
                    convertWavToAac(audioFile)
                }
                // AAC格式但缺少关键信息，尝试修复或重新编码
                audioMime == "audio/mp4a-latm" && (sampleRate <= 0 || channelCount <= 0) -> {
                    Log.d(TAG, "AAC file has invalid format info, trying to fix")
                    fixAacFormat(audioFile) ?: reEncodeAac(audioFile)
                }
                // 其他格式需要转换
                audioMime != "audio/mp4a-latm" -> {
                    Log.d(TAG, "Converting $audioMime to AAC format")
                    convertToAac(audioFile)
                }
                else -> {
                    Log.d(TAG, "Audio format is valid, using original file")
                    audioFile
                }
            }
            
            if (processedAudioFile == null) {
                Log.e(TAG, "Failed to process audio file")
                audioExtractor.release()
                return@withContext false
            }
            
            // 创建MediaMuxer
            val muxer = MediaMuxer(outputFile.absolutePath, MediaMuxer.OutputFormat.MUXER_OUTPUT_MPEG_4)
            
            // 提取视频轨道
            val videoExtractor = MediaExtractor()
            videoExtractor.setDataSource(videoFile.absolutePath)
            val videoTrackIndex = getVideoTrackIndex(videoExtractor)
            val videoFormat = videoExtractor.getTrackFormat(videoTrackIndex)
            val videoMuxerIndex = muxer.addTrack(videoFormat)
            
            // 提取处理后的音频轨道
            val processedAudioExtractor = MediaExtractor()
            processedAudioExtractor.setDataSource(processedAudioFile.absolutePath)
            val processedAudioTrackIndex = getAudioTrackIndex(processedAudioExtractor)
            val processedAudioFormat = processedAudioExtractor.getTrackFormat(processedAudioTrackIndex)
            
            // 验证音频格式信息
            val audioSampleRate = processedAudioFormat.getInteger(MediaFormat.KEY_SAMPLE_RATE)
            val audioChannels = processedAudioFormat.getInteger(MediaFormat.KEY_CHANNEL_COUNT)
            
            Log.d(TAG, "Processed audio - Sample rate: $audioSampleRate, Channels: $audioChannels")
            
            // 如果音频格式信息无效，创建标准格式
            val finalAudioFormat = if (audioSampleRate <= 0 || audioChannels <= 0) {
                Log.w(TAG, "Audio format info invalid, using default format")
                MediaFormat.createAudioFormat("audio/mp4a-latm", 44100, 2).apply {
                    setInteger(MediaFormat.KEY_BIT_RATE, 128000)
                    setInteger(MediaFormat.KEY_AAC_PROFILE, 2)
                }
            } else {
                processedAudioFormat
            }
            
            val audioMuxerIndex = muxer.addTrack(finalAudioFormat)
            
            // 开始混合
            muxer.start()
            
            // 复制视频数据
            videoExtractor.selectTrack(videoTrackIndex)
            copyMediaData(videoExtractor, muxer, videoMuxerIndex, false)
            
            // 复制音频数据
            processedAudioExtractor.selectTrack(processedAudioTrackIndex)
            copyMediaData(processedAudioExtractor, muxer, audioMuxerIndex, true)
            
            // 完成混合
            muxer.stop()
            muxer.release()
            videoExtractor.release()
            audioExtractor.release()
            processedAudioExtractor.release()
            
            // 清理临时文件
            if (processedAudioFile != audioFile) {
                processedAudioFile.delete()
                Log.d(TAG, "Cleaned up temporary AAC file")
            }
            
            Log.d(TAG, "Audio-video merge completed successfully")
            return@withContext true
            
        } catch (e: Exception) {
            Log.e(TAG, "Failed to merge audio and video with primary method", e)
            Log.d(TAG, "Trying fallback method...")
            
            // 尝试备用方法：仅复制视频文件
            return@withContext tryFallbackMerge(videoFile, audioFile, outputFile)
        }
    }
    
    /**
     * 备用合并方法：当主要方法失败时使用
     */
    private fun tryFallbackMerge(videoFile: File, audioFile: File?, outputFile: File): Boolean {
        try {
            Log.d(TAG, "Using fallback merge method")
            
            if (audioFile == null || !audioFile.exists()) {
                // 仅视频，直接复制
                Log.d(TAG, "No audio file, copying video only")
                videoFile.copyTo(outputFile, overwrite = true)
                return true
            }
            
            // 尝试简单的文件复制作为最后的备用方案
            Log.w(TAG, "Audio merge failed, saving video only")
            videoFile.copyTo(outputFile, overwrite = true)
            
            // 记录音频文件路径供后续处理
            Log.d(TAG, "Audio file preserved at: ${audioFile.absolutePath}")
            Log.d(TAG, "Video file saved at: ${outputFile.absolutePath}")
            
            return true
            
        } catch (e: Exception) {
            Log.e(TAG, "Fallback merge also failed", e)
            return false
        }
    }
    
    /**
     * 获取视频轨道索引
     */
    private fun getVideoTrackIndex(extractor: MediaExtractor): Int {
        for (i in 0 until extractor.trackCount) {
            val format = extractor.getTrackFormat(i)
            val mime = format.getString(MediaFormat.KEY_MIME)
            if (mime?.startsWith("video/") == true) {
                return i
            }
        }
        throw IOException("No video track found")
    }
    
    /**
     * 获取音频轨道索引
     */
    private fun getAudioTrackIndex(extractor: MediaExtractor): Int {
        for (i in 0 until extractor.trackCount) {
            val format = extractor.getTrackFormat(i)
            val mime = format.getString(MediaFormat.KEY_MIME)
            if (mime?.startsWith("audio/") == true) {
                return i
            }
        }
        throw IOException("No audio track found")
    }
    
    /**
     * 复制媒体数据
     */
    private fun copyMediaData(
        extractor: MediaExtractor,
        muxer: MediaMuxer,
        muxerIndex: Int,
        isAudio: Boolean
    ) {
        val bufferInfo = android.media.MediaCodec.BufferInfo()
        val buffer = java.nio.ByteBuffer.allocate(1024 * 1024) // 1MB buffer
        
        while (true) {
            val sampleSize = extractor.readSampleData(buffer, 0)
            if (sampleSize < 0) break
            
            bufferInfo.offset = 0
            bufferInfo.size = sampleSize
            bufferInfo.presentationTimeUs = extractor.sampleTime
            bufferInfo.flags = extractor.sampleFlags
            
            muxer.writeSampleData(muxerIndex, buffer, bufferInfo)
            extractor.advance()
        }
    }
    
    /**
     * 尝试修复AAC文件格式问题
     * @param aacFile 原始AAC文件
     * @return 修复后的AAC文件，失败时返回null
     */
    private fun fixAacFormat(aacFile: File): File? {
        try {
            Log.d(TAG, "Attempting to fix AAC format")
            
            // 创建修复后的AAC文件
            val fixedFile = File(aacFile.parent, "fixed_${System.currentTimeMillis()}.aac")
            
            // 使用MediaMuxer重新包装AAC数据
            val muxer = MediaMuxer(fixedFile.absolutePath, MediaMuxer.OutputFormat.MUXER_OUTPUT_MPEG_4)
            
            // 创建标准的AAC格式
            val format = MediaFormat.createAudioFormat("audio/mp4a-latm", 44100, 2)
            format.setInteger(MediaFormat.KEY_BIT_RATE, 128000)
            format.setInteger(MediaFormat.KEY_AAC_PROFILE, 2)
            format.setInteger(MediaFormat.KEY_MAX_INPUT_SIZE, 16384)
            
            val trackIndex = muxer.addTrack(format)
            muxer.start()
            
            // 读取原始AAC数据并重新写入
            val extractor = MediaExtractor()
            extractor.setDataSource(aacFile.absolutePath)
            val audioTrackIndex = getAudioTrackIndex(extractor)
            extractor.selectTrack(audioTrackIndex)
            
            val bufferInfo = MediaCodec.BufferInfo()
            val buffer = ByteBuffer.allocate(1024 * 1024)
            
            while (true) {
                val sampleSize = extractor.readSampleData(buffer, 0)
                if (sampleSize < 0) break
                
                bufferInfo.offset = 0
                bufferInfo.size = sampleSize
                bufferInfo.presentationTimeUs = extractor.sampleTime
                bufferInfo.flags = extractor.sampleFlags
                
                muxer.writeSampleData(trackIndex, buffer, bufferInfo)
                extractor.advance()
            }
            
            muxer.stop()
            muxer.release()
            extractor.release()
            
            if (fixedFile.exists() && fixedFile.length() > 0) {
                Log.d(TAG, "AAC format fixed successfully: ${fixedFile.absolutePath}")
                return fixedFile
            } else {
                Log.w(TAG, "AAC format fix failed")
                fixedFile.delete()
                return null
            }
            
        } catch (e: Exception) {
            Log.e(TAG, "Error fixing AAC format", e)
            return null
        }
    }
    
    /**
     * 重新编码AAC文件以修复格式问题
     * @param aacFile 原始AAC文件
     * @return 重新编码后的AAC文件，失败时返回null
     */
    private fun reEncodeAac(aacFile: File): File? {
        try {
            Log.d(TAG, "Starting AAC re-encoding")
            
            // 创建临时AAC文件
            val outputFile = File(aacFile.parent, "reencoded_${System.currentTimeMillis()}.aac")
            
            // 创建解码器
            val decoder = MediaCodec.createDecoderByType("audio/mp4a-latm")
            val encoder = MediaCodec.createEncoderByType("audio/mp4a-latm")
            
            // 配置解码器
            val inputFormat = MediaFormat.createAudioFormat("audio/mp4a-latm", 44100, 2)
            decoder.configure(inputFormat, null, null, 0)
            decoder.start()
            
            // 配置编码器
            val outputFormat = MediaFormat.createAudioFormat("audio/mp4a-latm", 44100, 2)
            outputFormat.setInteger(MediaFormat.KEY_BIT_RATE, 128000)
            outputFormat.setInteger(MediaFormat.KEY_AAC_PROFILE, 2) // AAC LC profile
            outputFormat.setInteger(MediaFormat.KEY_MAX_INPUT_SIZE, 16384)
            
            encoder.configure(outputFormat, null, null, MediaCodec.CONFIGURE_FLAG_ENCODE)
            encoder.start()
            
            // 执行解码和重新编码
            val success = transcodeAudio(decoder, encoder, aacFile, outputFile)
            
            decoder.stop()
            decoder.release()
            encoder.stop()
            encoder.release()
            
            if (success && outputFile.exists()) {
                Log.d(TAG, "AAC re-encoding completed: ${outputFile.absolutePath}")
                return outputFile
            } else {
                Log.e(TAG, "AAC re-encoding failed")
                outputFile.delete()
                return null
            }
            
        } catch (e: Exception) {
            Log.e(TAG, "Error re-encoding AAC", e)
            return null
        }
    }
    
    /**
     * 将任意音频格式转换为AAC
     * @param inputFile 输入音频文件
     * @return 转换后的AAC文件，失败时返回null
     */
    private fun convertToAac(inputFile: File): File? {
        try {
            Log.d(TAG, "Starting audio format conversion to AAC")
            
            // 创建临时AAC文件
            val outputFile = File(inputFile.parent, "converted_${System.currentTimeMillis()}.aac")
            
            // 创建编码器
            val encoder = MediaCodec.createEncoderByType("audio/mp4a-latm")
            
            // 配置编码器
            val format = MediaFormat.createAudioFormat("audio/mp4a-latm", 44100, 2)
            format.setInteger(MediaFormat.KEY_BIT_RATE, 128000)
            format.setInteger(MediaFormat.KEY_AAC_PROFILE, 2) // AAC LC profile
            format.setInteger(MediaFormat.KEY_MAX_INPUT_SIZE, 16384)
            
            encoder.configure(format, null, null, MediaCodec.CONFIGURE_FLAG_ENCODE)
            encoder.start()
            
            // 读取音频数据并编码
            val audioData = readAudioFile(inputFile)
            if (audioData == null) {
                Log.e(TAG, "Failed to read audio file")
                encoder.release()
                return null
            }
            
            val success = encodeAudioData(encoder, audioData, outputFile)
            
            encoder.stop()
            encoder.release()
            
            if (success && outputFile.exists()) {
                Log.d(TAG, "Audio conversion completed: ${outputFile.absolutePath}")
                return outputFile
            } else {
                Log.e(TAG, "Audio conversion failed")
                outputFile.delete()
                return null
            }
            
        } catch (e: Exception) {
            Log.e(TAG, "Error converting audio to AAC", e)
            return null
        }
    }
    
    /**
     * 将WAV文件转换为AAC格式
     * @param wavFile WAV音频文件
     * @return 转换后的AAC文件，失败时返回null
     */
    private fun convertWavToAac(wavFile: File): File? {
        try {
            Log.d(TAG, "Starting WAV to AAC conversion")
            
            // 创建临时AAC文件
            val aacFile = File(wavFile.parent, "temp_${System.currentTimeMillis()}.aac")
            
            // 创建AAC编码器
            val encoder = MediaCodec.createEncoderByType("audio/mp4a-latm")
            
            // 配置编码器
            val format = MediaFormat.createAudioFormat("audio/mp4a-latm", 44100, 2)
            format.setInteger(MediaFormat.KEY_BIT_RATE, 128000)
            format.setInteger(MediaFormat.KEY_AAC_PROFILE, 2) // AAC LC profile
            format.setInteger(MediaFormat.KEY_MAX_INPUT_SIZE, 16384)
            
            encoder.configure(format, null, null, MediaCodec.CONFIGURE_FLAG_ENCODE)
            encoder.start()
            
            // 读取WAV文件
            val wavData = readWavFile(wavFile)
            if (wavData == null) {
                Log.e(TAG, "Failed to read WAV file")
                encoder.release()
                return null
            }
            
            // 编码音频数据
            val success = encodeAudioData(encoder, wavData, aacFile)
            
            encoder.stop()
            encoder.release()
            
            if (success && aacFile.exists()) {
                Log.d(TAG, "WAV to AAC conversion completed: ${aacFile.absolutePath}")
                return aacFile
            } else {
                Log.e(TAG, "WAV to AAC conversion failed")
                aacFile.delete()
                return null
            }
            
        } catch (e: Exception) {
            Log.e(TAG, "Error converting WAV to AAC", e)
            return null
        }
    }
    
    /**
     * 读取WAV文件，跳过文件头，返回PCM数据
     */
    private fun readWavFile(wavFile: File): ByteArray? {
        try {
            val inputStream = wavFile.inputStream()
            val header = ByteArray(44) // WAV文件头大小
            inputStream.read(header)
            
            // 验证WAV文件头
            val riff = String(header, 0, 4)
            val wave = String(header, 8, 4)
            if (riff != "RIFF" || wave != "WAVE") {
                Log.e(TAG, "Invalid WAV file format")
                inputStream.close()
                return null
            }
            
            // 读取音频数据
            val audioData = inputStream.readBytes()
            inputStream.close()
            
            Log.d(TAG, "Read ${audioData.size} bytes of audio data from WAV file")
            return audioData
            
        } catch (e: Exception) {
            Log.e(TAG, "Error reading WAV file", e)
            return null
        }
    }
    
    /**
     * 使用MediaCodec编码音频数据
     */
    private fun encodeAudioData(encoder: MediaCodec, audioData: ByteArray, outputFile: File): Boolean {
        try {
            val outputStream = outputFile.outputStream()
            val bufferInfo = MediaCodec.BufferInfo()
            var inputBufferIndex: Int
            var outputBufferIndex: Int
            var inputOffset = 0
            var isEOS = false
            
            while (!isEOS) {
                // 输入数据
                inputBufferIndex = encoder.dequeueInputBuffer(10000)
                if (inputBufferIndex >= 0) {
                    val inputBuffer = encoder.getInputBuffer(inputBufferIndex)
                    inputBuffer?.clear()
                    
                    val chunkSize = minOf(4096, audioData.size - inputOffset)
                    if (chunkSize > 0) {
                        inputBuffer?.put(audioData, inputOffset, chunkSize)
                        encoder.queueInputBuffer(inputBufferIndex, 0, chunkSize, 0, 0)
                        inputOffset += chunkSize
                    } else {
                        // 结束输入
                        encoder.queueInputBuffer(inputBufferIndex, 0, 0, 0, MediaCodec.BUFFER_FLAG_END_OF_STREAM)
                        isEOS = true
                    }
                }
                
                // 输出数据
                outputBufferIndex = encoder.dequeueOutputBuffer(bufferInfo, 10000)
                if (outputBufferIndex >= 0) {
                    val outputBuffer = encoder.getOutputBuffer(outputBufferIndex)
                    if (outputBuffer != null && bufferInfo.size > 0) {
                        val outputData = ByteArray(bufferInfo.size)
                        outputBuffer.get(outputData)
                        outputStream.write(outputData)
                    }
                    
                    encoder.releaseOutputBuffer(outputBufferIndex, false)
                    
                    if ((bufferInfo.flags and MediaCodec.BUFFER_FLAG_END_OF_STREAM) != 0) {
                        isEOS = true
                    }
                }
            }
            
            outputStream.close()
            return true
            
        } catch (e: Exception) {
            Log.e(TAG, "Error encoding audio data", e)
            return false
        }
    }
    
    /**
     * 音频转码（解码+编码）
     */
    private fun transcodeAudio(
        decoder: MediaCodec,
        encoder: MediaCodec,
        inputFile: File,
        outputFile: File
    ): Boolean {
        try {
            val outputStream = outputFile.outputStream()
            val bufferInfo = MediaCodec.BufferInfo()
            var inputBufferIndex: Int
            var outputBufferIndex: Int
            var isEOS = false
            
            // 读取输入文件
            val inputData = inputFile.readBytes()
            var inputOffset = 0
            
            while (!isEOS) {
                // 解码输入
                inputBufferIndex = decoder.dequeueInputBuffer(10000)
                if (inputBufferIndex >= 0) {
                    val inputBuffer = decoder.getInputBuffer(inputBufferIndex)
                    inputBuffer?.clear()
                    
                    val chunkSize = minOf(4096, inputData.size - inputOffset)
                    if (chunkSize > 0) {
                        inputBuffer?.put(inputData, inputOffset, chunkSize)
                        decoder.queueInputBuffer(inputBufferIndex, 0, chunkSize, 0, 0)
                        inputOffset += chunkSize
                    } else {
                        decoder.queueInputBuffer(inputBufferIndex, 0, 0, 0, MediaCodec.BUFFER_FLAG_END_OF_STREAM)
                        isEOS = true
                    }
                }
                
                // 解码输出
                outputBufferIndex = decoder.dequeueOutputBuffer(bufferInfo, 10000)
                if (outputBufferIndex >= 0) {
                    val outputBuffer = decoder.getOutputBuffer(outputBufferIndex)
                    if (outputBuffer != null && bufferInfo.size > 0) {
                        // 将解码后的数据编码
                        val decodedData = ByteArray(bufferInfo.size)
                        outputBuffer.get(decodedData)
                        encodeAudioData(encoder, decodedData, outputFile)
                    }
                    
                    decoder.releaseOutputBuffer(outputBufferIndex, false)
                    
                    if ((bufferInfo.flags and MediaCodec.BUFFER_FLAG_END_OF_STREAM) != 0) {
                        isEOS = true
                    }
                }
            }
            
            outputStream.close()
            return true
            
        } catch (e: Exception) {
            Log.e(TAG, "Error transcoding audio", e)
            return false
        }
    }
    
    /**
     * 读取音频文件数据
     */
    private fun readAudioFile(audioFile: File): ByteArray? {
        return try {
            audioFile.readBytes()
        } catch (e: Exception) {
            Log.e(TAG, "Error reading audio file", e)
            null
        }
    }
    
    /**
     * 清理临时文件
     */
    fun cleanupTempFiles(vararg files: File) {
        files.forEach { file ->
            if (file.exists()) {
                try {
                    file.delete()
                    Log.d(TAG, "Deleted temp file: ${file.absolutePath}")
                } catch (e: Exception) {
                    Log.w(TAG, "Failed to delete temp file: ${file.absolutePath}", e)
                }
            }
        }
    }
}
